# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import gymnasium as gym
import torch
import numpy as np
import isaaclab.sim as sim_utils
from isaaclab.assets import Articulation, ArticulationCfg
from isaaclab.envs import DirectRLEnv, DirectRLEnvCfg
from isaaclab.envs.ui import BaseEnvWindow
from isaaclab.markers import VisualizationMarkers
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.sim import SimulationCfg
from isaaclab.terrains import TerrainImporterCfg
from isaaclab.utils import configclass
from isaaclab.utils.math import random_orientation, random_yaw_orientation, subtract_frame_transforms, yaw_angle, compute_pose_error, wrap_to_pi, euler_xyz_from_quat
from isaaclab.sensors import TiledCamera, TiledCameraCfg, ContactSensor, ContactSensorCfg

##
# Pre-defined configs
##
from isaaclab_assets import SCOUT_MINI_CFG  # isort: skip
from isaaclab.markers import CUBOID_MARKER_CFG, Sphere_MARKER_CFG  # isort: skip
from isaaclab.markers.config import BLUE_ARROW_X_MARKER_CFG, GREEN_ARROW_X_MARKER_CFG, RED_ARROW_X_MARKER_CFG, ROBOT_MARKER_CFG
from .differential_controller import BatchDifferentialController, Gen2DTrajectories

class ScoutMiniEnvWindow(BaseEnvWindow):
    """Window manager for the ScoutMini environment."""

    def __init__(self, env: ScoutMiniTrackingEnv, window_name: str = "IsaacLab"):
        """Initialize the window.

        Args:
            env: The environment object.
            window_name: The name of the window. Defaults to "IsaacLab".
        """
        # initialize base window
        super().__init__(env, window_name)
        # add custom UI elements
        with self.ui_window_elements["main_vstack"]:
            with self.ui_window_elements["debug_frame"]:
                with self.ui_window_elements["debug_vstack"]:
                    # add command manager visualization
                    self._create_debug_vis_ui_element("targets", self.env)

@configclass
class ScoutMiniTrackingEnvCfg(DirectRLEnvCfg):
    # env
    episode_length_s = 100.0
    decimation = 2
    action_space = 2
    observation_space = 16
    state_space = 0
    debug_vis = True

    ui_window_class_type = ScoutMiniEnvWindow

    # simulation
    sim: SimulationCfg = SimulationCfg(
        dt=1 / 100,
        render_interval=decimation,
        enable_scene_query_support = True,
        use_fabric=True, 
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=0.0,
        ),
    )
    terrain = TerrainImporterCfg(
        prim_path="/World/ground",
        terrain_type="plane",
        collision_group=-1,
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=0.0,
        ),
        debug_vis=False,
    )

    # scene
    scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing = 10.0, replicate_physics=True)
    # robot
    robot: ArticulationCfg = SCOUT_MINI_CFG.replace(prim_path="/World/envs/env_.*/Robot")

    action_scale = 1.0
    distance_threshold = 3.0

    # reward scales
    lin_vel_reward_scale = -0.05
    ang_vel_reward_scale = -0.05
    distance_to_goal_reward_scale = 15.0
    collision_reward_scale = -1.0
    action_rate_reward_scale = -0.1
    goal_reached_reward_scale = 0.5
    delta_goal_reward_scale = 200.0
    angle_reward_scale = 4.0
    target_angle_reward_scale = 2.0

class ScoutMiniTrackingEnv(DirectRLEnv):
    cfg: ScoutMiniTrackingEnvCfg

    def __init__(self, cfg: ScoutMiniTrackingEnvCfg, render_mode: str | None = None, **kwargs):
        super().__init__(cfg, render_mode, **kwargs)

        # Total thrust and moment applied to the base of the uav
        self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)
        self._previous_actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)

        # Goal position
        self._desired_pos_w = torch.zeros(self.num_envs, 3, device=self.device)
        self._desired_quat_w = torch.zeros(self.num_envs, 4, device=self.device)
        self._collision = torch.zeros(self.num_envs, device=self.device)
        self._is_goal_reached = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool)
        self._dist_to_goal = torch.zeros(self.num_envs, device=self.device)
        self._previous_dist_to_goal = torch.zeros(self.num_envs, device=self.device)
        self._path_max_count = torch.ones(self.num_envs, device=self.device)

        # Logging
        self._episode_sums = {
            key: torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
            for key in [
                "lin_vel",
                "ang_vel",
                "distance_to_goal",
                "collision",
                "action_rate_l2",
                "goal_reached",
                "delta_goal_dist",
                "angle_error",
                "reverse_penalty",
                "target_angle"
            ]
        }
        wheels_dof_names = ["front_left_wheel", "front_right_wheel", "rear_left_wheel", "rear_right_wheel"]

        # Get specific body indices
        self._wheels_dof_idx, _ = self._robot.find_joints(wheels_dof_names)
        self._robot_mass = self._robot.root_physx_view.get_masses()[0].sum()
        self._gravity_magnitude = torch.tensor(self.sim.cfg.gravity, device=self.device).norm()
        self._robot_weight = (self._robot_mass * self._gravity_magnitude).item()

        # 初始化控制器
        self.differential_controller = BatchDifferentialController(device=self.device)

        self.trajectories_controller = Gen2DTrajectories(num_objects=self.num_envs, device=self.device)
        # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
        self.set_debug_vis(self.cfg.debug_vis)

    def _setup_scene(self):
        self._robot = Articulation(self.cfg.robot)
        self.scene.articulations["robot"] = self._robot
        self.cfg.terrain.num_envs = self.scene.cfg.num_envs
        self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing
        self._terrain = self.cfg.terrain.class_type(self.cfg.terrain)
        # clone, filter, and replicate
        self.scene.clone_environments(copy_from_source=False)
        self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path])
        # add lights
        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
        light_cfg.func("/World/Light", light_cfg)

    def _pre_physics_step(self, actions: torch.Tensor):
        self._actions = self.cfg.action_scale * actions.clone().clamp(-1.0, 1.0)
        
    def _apply_action(self):
        left_and_right_wheel_speed = self.differential_controller.speed_to_angular(self._actions)
        left_and_right_wheel_speed[:, 0] *= -1 
        wheel_actions = torch.cat((left_and_right_wheel_speed, left_and_right_wheel_speed), dim=1)
        self._robot.set_joint_velocity_target(wheel_actions, self._wheels_dof_idx)

    def _get_observations(self) -> dict:
        self._previous_actions = self._actions.clone()
        self.update_desired_pos()
        pos_error, axis_angle_error = compute_pose_error(self._desired_pos_w, self._desired_quat_w, self._robot.data.root_pos_w, self._robot.data.root_quat_w)
        yaw_robo = yaw_angle(self._robot.data.root_quat_w)
        yaw_desired = yaw_angle(self._desired_quat_w)

        obs = torch.cat(
            [
                self._robot.data.root_lin_vel_b,
                self._robot.data.root_ang_vel_b,
                # self._robot.data.projected_gravity_b,
                pos_error,
                yaw_robo.unsqueeze(1),
                yaw_desired.unsqueeze(1),
                axis_angle_error[:, -1:],
                self._dist_to_goal.unsqueeze(1),
                self._actions,
            ],
            dim=-1,
        )
        observations = {"policy": obs}
        # print("obs:", observations)
        return observations

    def _get_rewards(self) -> torch.Tensor:
        lin_vel = torch.sum(torch.square(self._robot.data.root_lin_vel_b), dim=1)
        ang_vel = torch.sum(torch.square(self._robot.data.root_ang_vel_b), dim=1)
        self._dist_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w, dim=1)
        distance_to_goal_mapped = 1 - torch.tanh(0.5 * self._dist_to_goal)
        delta_goal_dist = self._previous_dist_to_goal - self._dist_to_goal
        # print("delta_goal_dist:", delta_goal_dist)
        self._previous_dist_to_goal = self._dist_to_goal.clone()

        # action rate
        action_rate = torch.sum(torch.square(self._actions - self._previous_actions), dim=1)
        self._is_goal_reached = self._dist_to_goal < 0.05
        yaw_robo = yaw_angle(self._robot.data.root_quat_w)
        # yaw_desired = yaw_angle(self._desired_quat_w)
        # angle_error = wrap_to_pi(yaw_robo - yaw_desired)
        pos_error, axis_angle_error = compute_pose_error(self._desired_pos_w, self._desired_quat_w, self._robot.data.root_pos_w, self._robot.data.root_quat_w)
        # print("pos_error:", pos_error, "axis_angle_error:", axis_angle_error, axis_angle_error[:,2], "angle_error_mapped:", torch.cos(axis_angle_error[:,2]))
        target_angle = torch.atan2(pos_error[:, 1], pos_error[:, 0])
        # print("target_angle:", target_angle)
        angel_error_to_target = yaw_robo - target_angle + torch.pi
        # print("distance_to_goal_mapped:", distance_to_goal_mapped)
        # print("angel_error_to_target:", angel_error_to_target, torch.cos(angel_error_to_target))
        angle_error_mapped = torch.cos(axis_angle_error[:,2])
        # target_angle_reward = torch.cos(angel_error_to_target) * torch.tanh(self._dist_to_goal)
        # print("target_angle_reward:",  target_angle_reward)
        # reverse_penalty = (self._robot.data.root_lin_vel_b[:, 0] < 0).float() * (-1.0)

        rewards = {
            "lin_vel": lin_vel * self.cfg.lin_vel_reward_scale * self.step_dt,
            "ang_vel": ang_vel * self.cfg.ang_vel_reward_scale * self.step_dt,
            "distance_to_goal": distance_to_goal_mapped * self.cfg.distance_to_goal_reward_scale * self.step_dt,
            "action_rate_l2": action_rate * self.cfg.action_rate_reward_scale * self.step_dt,
            "goal_reached": self._is_goal_reached.float() * self.cfg.goal_reached_reward_scale,
            "delta_goal_dist": delta_goal_dist * self.cfg.delta_goal_reward_scale * self.step_dt,
            "angle_error":  angle_error_mapped * distance_to_goal_mapped* self.cfg.angle_reward_scale * self.step_dt,
            # "reverse_penalty": reverse_penalty * self.step_dt,
            # "target_angle": target_angle_reward * self.cfg.target_angle_reward_scale * self.step_dt,
        }
        reward = torch.sum(torch.stack(list(rewards.values())), dim=0)
        # Logging
        for key, value in rewards.items():
            self._episode_sums[key] += value
        return reward

    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
        time_out = self.episode_length_buf >= self.max_episode_length - 1
        # time_out = torch.logical_or(self.episode_length_buf >= self.max_episode_length - 1, self._is_goal_reached)
        died = self._dist_to_goal > self.cfg.distance_threshold
        # died = self._robot.data.root_pos_w[:, 2] > 1.0
        # print("max_net_contact_forces:", max_net_contact_forces)
        # died = torch.logical_or(self._robot.data.root_pos_w[:, 2] < 0.5, self._robot.data.root_pos_w[:, 2] > 15.0)
        return died, time_out
    
    def _reset_idx(self, env_ids: torch.Tensor | None):
        if env_ids is None or len(env_ids) == self.num_envs:
            env_ids = self._robot._ALL_INDICES

        # Logging
        final_distance_to_goal = torch.linalg.norm(
            self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1
        ).mean()
        extras = dict()
        for key in self._episode_sums.keys():
            episodic_sum_avg = torch.mean(self._episode_sums[key][env_ids])
            extras["Episode_Reward/" + key] = episodic_sum_avg / self.max_episode_length_s
            self._episode_sums[key][env_ids] = 0.0
        self.extras["log"] = dict()
        self.extras["log"].update(extras)
        extras = dict()
        extras["Episode_Termination/died"] = torch.count_nonzero(self.reset_terminated[env_ids]).item()
        extras["Episode_Termination/time_out"] = torch.count_nonzero(self.reset_time_outs[env_ids]).item()
        extras["Metrics/final_distance_to_goal"] = final_distance_to_goal.item()
        self.extras["log"].update(extras)

        self._robot.reset(env_ids)
        super()._reset_idx(env_ids)
        self._collision[env_ids] = 0
        if len(env_ids) == self.num_envs:
            # Spread out the resets to avoid spikes in training when many environments reset at a similar time
            self.episode_length_buf = torch.randint_like(self.episode_length_buf, high=int(self.max_episode_length))

        self._actions[env_ids] = 0.0
        self._previous_actions[env_ids] = 0.0
        self._is_goal_reached[env_ids] &= False
        # Sample new commands
        self.update_desired_pos()

        # Reset robot state
        joint_pos = self._robot.data.default_joint_pos[env_ids]
        joint_vel = self._robot.data.default_joint_vel[env_ids]
        default_root_state = self._robot.data.default_root_state[env_ids]
        default_root_state[:, :3] = self._desired_pos_w[env_ids]
        default_root_state[:, 3:7] = self._desired_quat_w[env_ids]
        self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
        self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
        self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
        self._previous_dist_to_goal[env_ids] = torch.linalg.norm(self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1)

    def _set_debug_vis_impl(self, debug_vis: bool):
        # set visibility of markers
        # note: parent only deals with callbacks. not their visibility
        if debug_vis:
            # create markers if necessary for the first tome
            if not hasattr(self, "goal_pos_visualizer"):
                # -- goal
                marker_cfg = GREEN_ARROW_X_MARKER_CFG.copy()
                marker_cfg.prim_path = "/Visuals/Command/goal_position"
                marker_cfg.markers["arrow"].scale = (0.1, 0.1, 0.3)
                self.base_target_goal_visualizer = VisualizationMarkers(marker_cfg)
                # -- current
                marker_cfg = BLUE_ARROW_X_MARKER_CFG.copy()
                marker_cfg.prim_path = "/Visuals/Command/robot_current"
                marker_cfg.markers["arrow"].scale = (0.04, 0.04, 0.14)
                self.base_robot_visualizer = VisualizationMarkers(marker_cfg)
            # set their visibility to true
            self.base_target_goal_visualizer.set_visibility(True)
            self.base_robot_visualizer.set_visibility(True)
        else:
            if hasattr(self, "goal_pos_visualizer"):
                self.base_target_goal_visualizer.set_visibility(False)
                self.base_robot_visualizer.set_visibility(False)

    def _debug_vis_callback(self, event):
        # get marker location
        # -- base state
        robot_pos_w = self._robot.data.root_pos_w.clone()
        robot_quat_w = self._robot.data.root_quat_w.clone()
        robot_pos_w[:, 2] += 0.05
        # -- resolve the scales and quaternions
        # vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
        # vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
        # display markers
        self.base_target_goal_visualizer.visualize(self._desired_pos_w, self._desired_quat_w)
        self.base_robot_visualizer.visualize(robot_pos_w, robot_quat_w)

    def update_desired_pos(self):
        self._desired_pos_w, self._desired_quat_w = self.trajectories_controller.update()

@configclass
class ScoutMiniTrackingPlayEnvCfg(DirectRLEnvCfg):
    # env
    episode_length_s = 200.0
    decimation = 2
    action_space = 2
    observation_space = 16
    state_space = 0
    debug_vis = True

    ui_window_class_type = ScoutMiniEnvWindow

    # simulation
    sim: SimulationCfg = SimulationCfg(
        dt=1 / 100,
        render_interval=decimation,
        enable_scene_query_support = True,
        use_fabric=True, 
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=0.0,
        ),
    )
    terrain = TerrainImporterCfg(
        prim_path="/World/ground",
        terrain_type="plane",
        collision_group=-1,
        physics_material=sim_utils.RigidBodyMaterialCfg(
            friction_combine_mode="multiply",
            restitution_combine_mode="multiply",
            static_friction=1.0,
            dynamic_friction=1.0,
            restitution=0.0,
        ),
        debug_vis=False,
    )

    # scene
    scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=1, env_spacing = 10.0, replicate_physics=True)
    # robot
    robot: ArticulationCfg = SCOUT_MINI_CFG.replace(prim_path="/World/envs/env_.*/Robot")

    action_scale = 1.0
    distance_threshold = 0.8

class ScoutMiniTrackingPlayEnv(DirectRLEnv):
    cfg: ScoutMiniTrackingPlayEnvCfg

    def __init__(self, cfg: ScoutMiniTrackingPlayEnvCfg, render_mode: str | None = None, **kwargs):
        super().__init__(cfg, render_mode, **kwargs)

        # Total thrust and moment applied to the base of the uav
        self._actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)
        self._previous_actions = torch.zeros(self.num_envs, gym.spaces.flatdim(self.single_action_space), device=self.device)

        # Goal position
        self._desired_pos_w = 0.2 * torch.ones(self.num_envs, 3, device=self.device)
        self._desired_quat_w = random_yaw_orientation(self.num_envs, device=self.device)
        self._collision = torch.zeros(self.num_envs, device=self.device)
        self._is_goal_reached = torch.zeros(self.num_envs, device=self.device, dtype=torch.bool)
        self._dist_to_goal = torch.zeros(self.num_envs, device=self.device)
        self._previous_dist_to_goal = torch.zeros(self.num_envs, device=self.device)
        self._path_max_count = torch.ones(self.num_envs, device=self.device)

        # Logging

        wheels_dof_names = ["front_left_wheel", "front_right_wheel", "rear_left_wheel", "rear_right_wheel"]

        # Get specific body indices
        self._wheels_dof_idx, _ = self._robot.find_joints(wheels_dof_names)
        self._robot_mass = self._robot.root_physx_view.get_masses()[0].sum()
        self._gravity_magnitude = torch.tensor(self.sim.cfg.gravity, device=self.device).norm()
        self._robot_weight = (self._robot_mass * self._gravity_magnitude).item()

        # 初始化控制器
        self.differential_controller = BatchDifferentialController(device=self.device)

        # add handle for debug visualization (this is set to a valid handle inside set_debug_vis)
        self.set_debug_vis(self.cfg.debug_vis)

    def _setup_scene(self):
        self._robot = Articulation(self.cfg.robot)
        self.scene.articulations["robot"] = self._robot
        self.cfg.terrain.num_envs = self.scene.cfg.num_envs
        self.cfg.terrain.env_spacing = self.scene.cfg.env_spacing
        self._terrain = self.cfg.terrain.class_type(self.cfg.terrain)
        # clone, filter, and replicate
        self.scene.clone_environments(copy_from_source=False)
        self.scene.filter_collisions(global_prim_paths=[self.cfg.terrain.prim_path])
        # add lights
        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
        light_cfg.func("/World/Light", light_cfg)

    def _pre_physics_step(self, actions: torch.Tensor):
        self._actions = self.cfg.action_scale * actions.clone().clamp(-1.0, 1.0)
        
    def _apply_action(self):
        left_and_right_wheel_speed = self.differential_controller.speed_to_angular(self._actions)
        left_and_right_wheel_speed[:, 0] *= -1 
        wheel_actions = torch.cat((left_and_right_wheel_speed, left_and_right_wheel_speed), dim=1)
        self._robot.set_joint_velocity_target(wheel_actions, self._wheels_dof_idx)
    
    def _get_observations(self) -> dict:
        self._previous_actions = self._actions.clone()
        pos_error, axis_angle_error = compute_pose_error(self._desired_pos_w, self._desired_quat_w, self._robot.data.root_pos_w, self._robot.data.root_quat_w)
        yaw_robo = yaw_angle(self._robot.data.root_quat_w)
        yaw_desired = yaw_angle(self._desired_quat_w)

        obs = torch.cat(
            [
                self._robot.data.root_lin_vel_b,
                self._robot.data.root_ang_vel_b,
                # self._robot.data.projected_gravity_b,
                pos_error,
                yaw_robo.unsqueeze(1),
                yaw_desired.unsqueeze(1),
                axis_angle_error[:, -1:],
                self._dist_to_goal.unsqueeze(1),
                self._actions,
            ],
            dim=-1,
        )
        observations = {"policy": obs}
        # print("obs:", observations)
        return observations

    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
        time_out = self.episode_length_buf >= self.max_episode_length - 1
        died = self._dist_to_goal > self.cfg.distance_threshold
        return died, time_out
    
    def _reset_idx(self, env_ids: torch.Tensor | None):
        if env_ids is None or len(env_ids) == self.num_envs:
            env_ids = self._robot._ALL_INDICES

        self._robot.reset(env_ids)
        super()._reset_idx(env_ids)
        self._collision[env_ids] = 0
        if len(env_ids) == self.num_envs:
            # Spread out the resets to avoid spikes in training when many environments reset at a similar time
            self.episode_length_buf = torch.randint_like(self.episode_length_buf, high=int(self.max_episode_length))

        self._actions[env_ids] = 0.0
        self._previous_actions[env_ids] = 0.0
        self._is_goal_reached[env_ids] &= False
        # Sample new commands
        # Reset robot state
        joint_pos = self._robot.data.default_joint_pos[env_ids]
        joint_vel = self._robot.data.default_joint_vel[env_ids]
        default_root_state = self._robot.data.default_root_state[env_ids]
        default_root_state[:, :3] = self._desired_pos_w[env_ids]
        default_root_state[:, 3:7] = self._desired_quat_w[env_ids]
        self._robot.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
        self._robot.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
        self._robot.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
        self._previous_dist_to_goal[env_ids] = torch.linalg.norm(self._desired_pos_w[env_ids] - self._robot.data.root_pos_w[env_ids], dim=1)

    def _set_debug_vis_impl(self, debug_vis: bool):
        # set visibility of markers
        # note: parent only deals with callbacks. not their visibility
        if debug_vis:
            # create markers if necessary for the first tome
            if not hasattr(self, "goal_pos_visualizer"):
                # -- goal
                marker_cfg = ROBOT_MARKER_CFG.copy()
                marker_cfg.prim_path = "/Visuals/Command/goal_position"
                self.base_target_goal_visualizer = VisualizationMarkers(marker_cfg)
                # -- current
                marker_cfg = BLUE_ARROW_X_MARKER_CFG.copy()
                marker_cfg.prim_path = "/Visuals/Command/robot_current"
                marker_cfg.markers["arrow"].scale = (0.04, 0.04, 0.14)
                self.base_robot_visualizer = VisualizationMarkers(marker_cfg)
            # set their visibility to true
            self.base_target_goal_visualizer.set_visibility(True)
            self.base_robot_visualizer.set_visibility(True)
        else:
            if hasattr(self, "goal_pos_visualizer"):
                self.base_target_goal_visualizer.set_visibility(False)
                self.base_robot_visualizer.set_visibility(False)

    def _debug_vis_callback(self, event):
        # get marker location
        # -- base state
        robot_pos_w = self._robot.data.root_pos_w.clone()
        robot_quat_w = self._robot.data.root_quat_w.clone()
        robot_pos_w[:, 2] += 0.1
        # -- resolve the scales and quaternions
        # vel_des_arrow_scale, vel_des_arrow_quat = self._resolve_xy_velocity_to_arrow(self.command[:, :2])
        # vel_arrow_scale, vel_arrow_quat = self._resolve_xy_velocity_to_arrow(self.robot.data.root_lin_vel_b[:, :2])
        # display markers
        self.base_target_goal_visualizer.visualize(self._desired_pos_w, self._desired_quat_w)
        self.base_robot_visualizer.visualize(robot_pos_w, robot_quat_w)

    def _get_rewards(self) -> torch.Tensor:
        self._dist_to_goal = torch.linalg.norm(self._desired_pos_w - self._robot.data.root_pos_w, dim=1)
        distance_to_goal_mapped = 1 - torch.tanh(self._dist_to_goal / 0.8)
        rewards = {
            "distance_to_goal": distance_to_goal_mapped,
        }
        reward = torch.sum(torch.stack(list(rewards.values())), dim=0)
        return reward

    def update_desired_pose(self, pos, quat):
        self._desired_pos_w = torch.from_numpy(pos).float().to(self.device)
        self._desired_quat_w = torch.from_numpy(quat).float().to(self.device)